package com.mrx;

import com.mrx.loader.AgentClassLoader;
import com.mrx.loader.protocol.ProtocolRegister;
import com.mrx.loader.util.EntryStream;
import org.apache.commons.compress.archivers.tar.TarArchiveInputStream;
import org.apache.commons.compress.compressors.gzip.GzipCompressorInputStream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.lang.reflect.Method;
import java.net.MalformedURLException;
import java.net.URL;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
import java.util.jar.Attributes;
import java.util.jar.JarEntry;
import java.util.jar.JarInputStream;
import java.util.stream.Collectors;

public class Main {

    private static final Logger logger = LoggerFactory.getLogger(Main.class);

    public static void main(String[] args) throws Exception {
        File packageFile = new File("HelloWorld/target/package/HelloWorld-1.0-SNAPSHOT.jar");
        File tarFile = new File("HelloWorld/target/HelloWorld-1.0-SNAPSHOT.tar.gz");
        ProtocolRegister.register();
        invokeJarWithLibsAsGzip(tarFile.toURI().toURL(), "HelloWorld-1.0-SNAPSHOT/HelloWorld-1.0-SNAPSHOT.jar");
        //        invokeJarWithLibs(packageFile.toURI().toURL());
    }

    /**
     * 执行打包成 tar.gz 的项目
     *
     * @param gzipUrl     gz 文件路径 URL
     * @param bootJarPath 启动 jar 包路径
     */
    @SuppressWarnings("SameParameterValue")
    private static void invokeJarWithLibsAsGzip(URL gzipUrl, String bootJarPath) throws IOException, ReflectiveOperationException {
        try (TarArchiveInputStream tarInputStream = new TarArchiveInputStream(new GzipCompressorInputStream(gzipUrl.openStream()))) {
            String urlString = gzipUrl.toString().replaceFirst("file", "gzip");
            // 创建一个 AgentClassLoader 并将其设置为 ContextClassLoader, SpringBoot 使用 ContextClassLoader 加载 class 和 resource
            AgentClassLoader loader = new AgentClassLoader();
            Thread.currentThread().setContextClassLoader(loader);
            Map<String, byte[]> mainJarMap = new HashMap<>();
            AtomicReference<String> mainClass = new AtomicReference<>();
            EntryStream.stream(() -> EntryStream.getStreamTarEntry(tarInputStream)).forEach(entry -> {
                // 读取 tar 文件里的 jar 文件, 将其组装成自定义协议的 URL 对象, 方便后期将其加入 classPath 使用
                String name = entry.getName();
                if (!entry.isDirectory() && name.endsWith(".jar")) {
                    URL lib;
                    try {
                        lib = new URL(urlString + "gzip://" + entry.getName());
                    } catch (MalformedURLException e) {
                        throw new RuntimeException(e);
                    }
                    // 对启动 jar 包做特殊操作, 它可能依赖别的库, 所以得最后加载
                    if (name.equals(bootJarPath)) {
                        try {
                            JarInputStream jarFile = new JarInputStream(new ByteArrayInputStream(tarInputStream.readAllBytes()));
                            mainClass.set(jarFile.getManifest().getMainAttributes().getValue(Attributes.Name.MAIN_CLASS));
                            EntryStream.stream(() -> EntryStream.getStreamJarEntry(jarFile)).filter(it -> it.getName().endsWith(".class")).forEach(it -> {
                                try {
                                    String clazzName = it.getName().replace("/", ".")
                                            .replace(".class", "");
                                    mainJarMap.put(clazzName, jarFile.readAllBytes());
                                } catch (IOException e) {
                                    throw new RuntimeException(e);
                                }
                            });
                        } catch (IOException e) {
                            throw new RuntimeException(e);
                        }
                    }
                    loader.addURL(lib);
                }
            });
            // 对启动 jar 进行翻译操作, 修改其字节码
            for (Map.Entry<String, byte[]> entry : mainJarMap.entrySet()) {
                loader.loadClass(entry.getKey(), entry.getValue());
            }
            // 加载启动类并执行其 main 方法
            Class<?> mainClazz = loader.loadClass(mainClass.get());
            Method main = mainClazz.getMethod("main", String[].class);
            main.invoke(null, (Object) new String[]{});
        }
    }

    /**
     * 执行打包成 jar 的项目
     *
     * @param packageFile jar 文件路径 URL
     */
    @SuppressWarnings("unused")
    private static void invokeJarWithLibs(URL packageFile) throws IOException, ReflectiveOperationException {
        String path = packageFile.toString();
        String parentPath = path.substring(0, path.lastIndexOf("/") + 1);
        // 读取 jar 包
        try (JarInputStream jarFile = new JarInputStream(packageFile.openStream())) {
            AgentClassLoader loader = new AgentClassLoader();
            Thread.currentThread().setContextClassLoader(loader);
            Attributes attributes = jarFile.getManifest().getMainAttributes();
            String mainClass = attributes.getValue(Attributes.Name.MAIN_CLASS);
            String classPath = attributes.getValue(Attributes.Name.CLASS_PATH);
            // 读取 manifest 中的 classPath, 加载额外 jar 包
            List<URL> libs = Arrays.stream(classPath.split(" "))
                    .filter(it -> !it.equals("."))
                    .map(parentPath::concat)
                    .map(Main::newURL)
                    .collect(Collectors.toList());
            // 将额外 jar 加入 classPath
            loader.addURLs(libs);
            // 将自己也加入 classPath, 不然会没法找到启动 jar 中的资源文件, 以及 class
            loader.addURL(packageFile);
            // 加载启动 jar 包里的所有 class 文件, 并对其执行翻译操作
            for (JarEntry entry; (entry = jarFile.getNextJarEntry()) != null; ) {
                if (entry.getName().endsWith(".class")) {
                    String clazzName = entry.getName().replace("/", ".")
                            .replace(".class", "");
                    loader.loadClass(clazzName, jarFile.readAllBytes());
                }
            }
            // 加载 main 方法并执行
            logger.debug("loading mainClass: {}", mainClass);
            Class<?> mainClazz = loader.loadClass(mainClass);
            Method main = mainClazz.getMethod("main", String[].class);
            main.invoke(null, (Object) new String[]{});
        }
    }

    private static URL newURL(String path) {
        try {
            return new URL(path);
        } catch (MalformedURLException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 执行 .class 文件
     */
    @SuppressWarnings("unused")
    private static void invokeMain() throws IOException, ReflectiveOperationException {
        final String CLASS_FILE = "HelloWorld/target/classes/com/mrx/example/HelloWorldApplication.class";
        final String MAIN_CLASS = "com.mrx.example.HelloWorldApplication";
        try (AgentClassLoader loader = new AgentClassLoader()) {
            Class<?> clazz = loader.loadClass(MAIN_CLASS, Files.readAllBytes(Path.of(CLASS_FILE)));
            Method main = clazz.getDeclaredMethod("main", String[].class);
            main.invoke(null, (Object) new String[]{});
        }
    }

}